%% On the Optimal Design of Transfers and Income-Tax Progressivity
% Discretization of productivity processes

clear; clc; close all;
addpath(genpath(pwd));

%% Parameters

% Set parameters of Gaussian Mixture Autoregressive Process
% [P,X] = discreteGMAR(mu,A,pC,muC,sigmaC,Nm,nMoments,method,nSigmas)
mu = 0;                 % unconditional mean
rho = 0.935;            % persistence of GMAR(1)
p1 = 0.85;              % probability of drawing from normal 1
p2 = 1-p1;              % probability of drawing from normal 2
pvec = [p1 p2];         % vector of probabilities
mu1 = 0.017;            % mean of normal 1
mu2 = -(p1/(1-p1))*mu1; % mean of normal 2
muvec = [mu1 mu2];      % vector of means
sig1 = 1.07*0.155;      % standard deviation of normal 1     
sig2 = 1.07*0.50;       % standard deviation of normal 2
sigvec = [sig1 sig2];   % vector of standard deviations
n = 19;                 % number of grid points
nmom = 2;               % number of moments to match (default = 2)
method = 'even';        % quadrature method (default = 'even')
space = 3;              % grid spacing when using even-spaced grid

%% GMAR(1) Discretization

% Call discretization
[pi_mix,level_mix] = discreteGMAR(mu,rho,pvec,muvec,sigvec,n,nmom,method,space);
level_mix = exp(level_mix)';

% Stationary distribution
x = ones(1,n)/n;                    % initialize stationary distribution
error_x = 1;                        % initialize error
while (error_x > 1e-10)
    xnew = x*pi_mix;                % apply transition matrix
    error_x = max(abs(xnew-x));     % check convergence
    x = xnew;                       % update guess for stationary distribution
end
dist_mix = (x/sum(x))';             % stationary distribution
cum_dist_mix = cumsum(dist_mix);    % cumulative stationary distribution

% Store results
% Transition matrix GMAR
fid=fopen('Px_GMAR.txt','w');
fprintf(fid, '%12.16f\n', pi_mix);
fclose(fid);

% Grid without Pareto
fid=fopen('x_vec_nopareto.txt','w');
fprintf(fid, '%12.16f\n', level_mix);
fclose(fid);

% Compute average productivity by quantile
% Note that first everything is expressed in (productivity/#quantiles) units; will be multiplied below
% This case will be stored in second column of x_quantile
numel_q = 10;                   % number of quantiles
x_quantile = ones(numel_q,3);   % to store average productivities by quantile for three cases: GMAR No Pareto; GMAR Pareto; Gmar Richer Poor   
iq = 1;                         % first quantile                           
a = 1;                          % index of lower bound of quantile 1
b = find(cum_dist_mix>iq/numel_q,1,'first');    % index of upper bound of quantile 1
x_quantile(1,2) = sum(dist_mix(a:b-1).*level_mix(a:b-1))+level_mix(b).*(iq/numel_q-cum_dist_mix(b-1));  % average productivity in quantile 1
a = b;                          % set upper bound of quantile 1 to lower bound of next quantile
for iq = 2:numel_q-1            % loop over quantiles 2 to to n-1
    b = find(cum_dist_mix>iq/numel_q,1,'first');    % index of upper bound of quantile
    if (b==a)                                       % if upper bound equals lower bound
        x_quantile(iq,2) = level_mix(b)/numel_q;    % avg. productivity is equal to productivity at that index
    else                                            % otherwise average
        x_quantile(iq,2) = sum(dist_mix(a+1:b-1).*level_mix(a+1:b-1)) + ...     % average for groups within a and b
        (cum_dist_mix(a)-(iq-1)/numel_q).*level_mix(a) + ...                    % appropriate share from a
        level_mix(b).*(iq/numel_q-cum_dist_mix(b-1));                           % appropriate share from b
    end
    a = b;                                          % reset upper bound to lower bound
end
x_quantile(end,2) = sum(dist_mix(a+1:end).*level_mix(a+1:end)) + (cum_dist_mix(a)-(iq)/numel_q).*level_mix(a);  % last quantile

%% Pareto Adjustment for GMAR(1)

% Call discretization (have to rerun because already exponentiated above)
[pi_mix,level_mix] = discreteGMAR(mu,rho,pvec,muvec,sigvec,n,nmom,method,space);

% Pareto parameters
cum_dist_rev = cumsum(dist_mix,'reverse');      % cumulative distribution from top to bottom 
cutoff = find(cum_dist_rev<0.15,1,'first');     % cutoff for starting Pareto tail
x_m = 1-cum_dist_rev(cutoff);                   % mass above cutoff
kappa = 1.6;                                    % Pareto tail parameter

% Inputs to Pareto quantile function (see Hubmer, Krusell, Smith)
fx = ones(n,1);
for ii = cutoff:n
    fx(ii) = 1-cum_dist_rev(ii);
end
arg_par = ones(n,1);
for ii = cutoff:n
    arg_par(ii) = (fx(ii)-x_m)/(1-x_m);
end
 
% Grid points according to Pareto (using quantile function)
x = ones(n,1);
for ii = cutoff:n
    x(ii) = exp(level_mix(cutoff))/((1-arg_par(ii))^(1/kappa));
end

% Generate complete grid
level_mix = exp(level_mix)';
xgrid_pareto = [level_mix(1:cutoff);x(cutoff+1:n)];

% Store results
fid=fopen('x_vec.txt','w');
fprintf(fid, '%12.16f\n', xgrid_pareto);
fclose(fid);

% Compute average productivity by quantile
% Note that first everything is expressed in (productivity/#quantiles) units; will be multiplied below
% This case will be stored in first column of x_quantile
iq = 1;                         % first quantile                           
a = 1;                          % index of lower bound of quantile 1
b = find(cum_dist_mix>iq/numel_q,1,'first');    % index of upper bound of quantile 1
x_quantile(1,1) = sum(dist_mix(a:b-1).*xgrid_pareto(a:b-1))+xgrid_pareto(b).*(iq/numel_q-cum_dist_mix(b-1));  % average productivity in quantile 1
a = b;                          % set upper bound of quantile 1 to lower bound of next quantile
for iq = 2:numel_q-1            % loop over quantiles 2 to to n-1
    b = find(cum_dist_mix>iq/numel_q,1,'first');    % index of upper bound of quantile
    if (b==a)                                       % if upper bound equals lower bound
        x_quantile(iq,1) = xgrid_pareto(b)/numel_q; % avg. productivity is equal to productivity at that index
    else                                            % otherwise average
        x_quantile(iq,1) = sum(dist_mix(a+1:b-1).*xgrid_pareto(a+1:b-1)) + ...  % average for groups within a and b
        (cum_dist_mix(a)-(iq-1)/numel_q).*xgrid_pareto(a) + ...                 % appropriate share from a
        xgrid_pareto(b).*(iq/numel_q-cum_dist_mix(b-1));                        % appropriate share from b
    end
    a = b;                                          % reset upper bound to lower bound
end
x_quantile(end,1) = sum(dist_mix(a+1:end).*xgrid_pareto(a+1:end)) + (cum_dist_mix(a)-(iq)/numel_q).*xgrid_pareto(a);  % last quantile

%% Richer poor

% Reset bottom income states below 20th percentile to value at 20th percentile
cutoff = find(cum_dist_mix>0.20,1,'first');     % cutoff until which productivities are changed
xgrid_pareto = [xgrid_pareto(cutoff)*ones(cutoff-1,1);xgrid_pareto(cutoff:end)];    % adjust productivity states

% Store results
fid=fopen('x_vec_richerpoor.txt','w');
fprintf(fid, '%12.16f\n', xgrid_pareto);
fclose(fid);

% Compute average productivity by quantile
% Note that first everything is expressed in (productivity/#quantiles) units; will be multiplied below
% This case will be stored in third column of x_quantile
iq = 1;                         % first quantile                           
a = 1;                          % index of lower bound of quantile 1
b = find(cum_dist_mix>iq/numel_q,1,'first');    % index of upper bound of quantile 1
x_quantile(1,3) = sum(dist_mix(a:b-1).*xgrid_pareto(a:b-1))+xgrid_pareto(b).*(iq/numel_q-cum_dist_mix(b-1));  % average productivity in quantile 1
a = b;                          % set upper bound of quantile 1 to lower bound of next quantile
for iq = 2:numel_q-1            % loop over quantiles 2 to to n-1
    b = find(cum_dist_mix>iq/numel_q,1,'first');    % index of upper bound of quantile
    if (b==a)                                       % if upper bound equals lower bound
        x_quantile(iq,3) = xgrid_pareto(b)/numel_q; % avg. productivity is equal to productivity at that index
    else                                            % otherwise average
        x_quantile(iq,3) = sum(dist_mix(a+1:b-1).*xgrid_pareto(a+1:b-1)) + ...  % average for groups within a and b
        (cum_dist_mix(a)-(iq-1)/numel_q).*xgrid_pareto(a) + ...                 % appropriate share from a
        xgrid_pareto(b).*(iq/numel_q-cum_dist_mix(b-1));                        % appropriate share from b
    end
    a = b;                                          % reset upper bound to lower bound
end
x_quantile(end,3) = sum(dist_mix(a+1:end).*xgrid_pareto(a+1:end)) + (cum_dist_mix(a)-(iq)/numel_q).*xgrid_pareto(a);  % last quantile

% Recover true average productivities for all cases
x_quantile = x_quantile*numel_q;

%% AR(1) Discretization

var_mix = p1*(mu1^2+sig1^2) + p2*(mu2^2+sig2^2);    % variance corresponding to mixture
std_mix = sqrt(var_mix);                            % standard deviation corresponding to mixture

[pi_norm,level_norm,dist_norm,alambda,asigmay]=markovappr(rho,sqrt(var_mix),space,n);   % Tauchen
level_norm = exp(level_norm)';                                                          % transform to levels

% Store results
fid=fopen('Px_AR.txt','w');
fprintf(fid, '%12.16f\n', pi_norm);
fclose(fid);


%% AR(1) Discretization; robustness with higher rho

rho_robust = 0.95; 
var_mix_robust = (var_mix/(1-rho^2))*(1-rho_robust^2);

[pi_norm_robust,level_norm_robust,dist_norm_robust,~,~]=markovappr(rho_robust,sqrt(var_mix_robust),space,n);    % Tauchen
level_norm_robust = exp(level_norm_robust)';                                                                    % transform to levels

fid=fopen('Px_AR_Robust.txt','w');
fprintf(fid, '%12.16f\n', pi_norm_robust);
fclose(fid);

%% Save data for figure

% Save average productivities by quantile for figure
filename='x_quantile.mat';
save(filename,'x_quantile')

%% GMAR Pareto for alternative preference parameters

% Set parameters of Gaussian Mixture Autoregressive Process
% [P,X] = discreteGMAR(mu,A,pC,muC,sigmaC,Nm,nMoments,method,nSigmas)
mu = 0;                 % unconditional mean
rho = 0.935;            % persistence of GMAR(1)
p1 = 0.85;              % probability of drawing from normal 1
p2 = 1-p1;              % probability of drawing from normal 2
pvec = [p1 p2];         % vector of probabilities
mu1 = 0.021;            % mean of normal 1
mu2 = -(p1/(1-p1))*mu1; % mean of normal 2
muvec = [mu1 mu2];      % vector of means
sig1 = 0.15;            % standard deviation of normal 1     
sig2 = 0.47;            % standard deviation of normal 2
sigvec = [sig1 sig2];   % vector of standard deviations
n = 19;                 % number of grid points
nmom = 2;               % number of moments to match (default = 2)
method = 'even';        % quadrature method (default = 'even')
space = 3;              % grid spacing when using even-spaced grid

% Call discretization
[pi_mix,level_mix] = discreteGMAR(mu,rho,pvec,muvec,sigvec,n,nmom,method,space);
level_mix = exp(level_mix)';

% Stationary distribution
x = ones(1,n)/n;                    % initialize stationary distribution
error_x = 1;                        % initialize error
while (error_x > 1e-10)
    xnew = x*pi_mix;                % apply transition matrix
    error_x = max(abs(xnew-x));     % check convergence
    x = xnew;                       % update guess for stationary distribution
end
dist_mix = (x/sum(x))';             % stationary distribution
cum_dist_mix = cumsum(dist_mix);    % cumulative stationary distribution

% Store transition matrix
fid=fopen('Px_Pref.txt','w');
fprintf(fid, '%12.16f\n', pi_mix);
fclose(fid);

% Call discretization (have to rerun because already exponentiated above)
[pi_mix,level_mix] = discreteGMAR(mu,rho,pvec,muvec,sigvec,n,nmom,method,space);

% Pareto parameters
cum_dist_rev = cumsum(dist_mix,'reverse');      % cumulative distribution from top to bottom 
cutoff = find(cum_dist_rev<0.15,1,'first');     % cutoff for starting Pareto tail
x_m = 1-cum_dist_rev(cutoff);                   % mass above cutoff
kappa = 1.6;                                    % Pareto tail parameter

% Inputs to Pareto quantile function (see Hubmer, Krusell, Smith)
fx = ones(n,1);
for ii = cutoff:n
    fx(ii) = 1-cum_dist_rev(ii);
end
arg_par = ones(n,1);
for ii = cutoff:n
    arg_par(ii) = (fx(ii)-x_m)/(1-x_m);
end
 
% Grid points according to Pareto (using quantile function)
x = ones(n,1);
for ii = cutoff:n
    x(ii) = exp(level_mix(cutoff))/((1-arg_par(ii))^(1/kappa));
end

% Generate complete grid
level_mix = exp(level_mix)';
xgrid_pareto = [level_mix(1:cutoff);x(cutoff+1:n)];

% Store grid
fid=fopen('x_vec_pref.txt','w');
fprintf(fid, '%12.16f\n', xgrid_pareto);
fclose(fid);


